import os, argparse, warnings
from pathlib import Path
from datetime import datetime

ROOT_DIR=Path(os.getcwd())
MODELS_DIR = ROOT_DIR / "models"
TEMPORARY_DIR = ROOT_DIR / "temp"

# 添加当前目录到环境变量
os.environ['PYTHONPATH'] = os.path.dirname(__file__)

# 配置 huggingface_hub 环境变量
os.environ['HF_ENDPOINT']='https://hf-mirror.com'
os.environ['HF_HOME'] = MODELS_DIR.as_posix()
os.environ['HF_HUB_CACHE'] = MODELS_DIR.as_posix()
os.environ['HF_HUB_DISABLE_SYMLINKS_WARNING'] = 'true'

# 抑制警告
warnings.filterwarnings("ignore")

from transcribe import AudioTranscribe


# 模型常量
DEFAULT_MODEL = "parakeet-tdt-0.6b-v3"
PARAKEET_TDT_06B_V2 = "parakeet-tdt-0.6b-v2"
JAPANESE_MODEL = "parakeet-tdt_ctc-0.6b-ja"

# 语言配置信息
language_config = {
    "bg": {"name": "保加利亚语", "model": DEFAULT_MODEL},
    "hr": {"name": "克罗地亚语", "model": DEFAULT_MODEL},
    "cs": {"name": "捷克语", "model": DEFAULT_MODEL},
    "da": {"name": "丹麦语", "model": DEFAULT_MODEL},
    "nl": {"name": "荷兰语", "model": DEFAULT_MODEL},
    "en": {"name": "英语", "model": PARAKEET_TDT_06B_V2},
    "et": {"name": "爱沙尼亚语", "model": DEFAULT_MODEL},
    "fi": {"name": "芬兰语", "model": DEFAULT_MODEL},
    "fr": {"name": "法语", "model": DEFAULT_MODEL},
    "de": {"name": "德语", "model": DEFAULT_MODEL},
    "el": {"name": "希腊语", "model": DEFAULT_MODEL},
    "hu": {"name": "匈牙利语", "model": DEFAULT_MODEL},
    "it": {"name": "意大利语", "model": DEFAULT_MODEL},
    "lv": {"name": "拉脱维亚语", "model": DEFAULT_MODEL},
    "lt": {"name": "立陶宛语", "model": DEFAULT_MODEL},
    "mt": {"name": "马耳他语", "model": DEFAULT_MODEL},
    "pl": {"name": "波兰语", "model": DEFAULT_MODEL},
    "pt": {"name": "葡萄牙语", "model": DEFAULT_MODEL},
    "ro": {"name": "罗马尼亚语", "model": DEFAULT_MODEL},
    "sk": {"name": "斯洛伐克语", "model": DEFAULT_MODEL},
    "sl": {"name": "斯洛文尼亚语", "model": DEFAULT_MODEL},
    "es": {"name": "西班牙语", "model": DEFAULT_MODEL},
    "sv": {"name": "瑞典语", "model": DEFAULT_MODEL},
    "ru": {"name": "俄语", "model": DEFAULT_MODEL},
    "uk": {"name": "乌克兰语", "model": DEFAULT_MODEL},
    "ja": {"name": "日语", "model": JAPANESE_MODEL}
}

# 模型列表
model_list = {code: info["model"] for code, info in language_config.items()}

# 语言列表
languages = list(model_list.keys())

def parse_arguments():
    """
    解析命令行参数
    """
    parser = argparse.ArgumentParser(description='音频转录工具')
    parser.add_argument('-i', '--input', required=True, help='输入音频文件路径')
    parser.add_argument('-o', '--output', required=True, help='输出字幕文件路径')
    parser.add_argument('-l', '--language', required=False, choices=['en', 'ja'], help='语言代码: en(英语) 或 ja(日语)')
    parser.add_argument('-t', '--task_id', required=False, help='任务ID')
    parser.add_argument('-c', '--chunk_minutes', type=int, default=5, help='分块分钟数，默认5分钟')
    parser.add_argument('-s', '--min_silence_duration', type=float, default=0.2, help='最小静音时长，默认0.2秒')
    parser.add_argument('-m', '--model', required=False, help='模型名称，默认为 parakeet-tdt-0.6b-v3')
    parser.add_argument("--temp_dir", required=False, help="临时目录，默认为当前目录下的 temp 目录")
    
    return parser.parse_args()


# 音频转录
def transcribe_audio():
    args = parse_arguments()

    if(args.task_id == None):
        now = datetime.now()
        formatted_time = now.strftime("%Y%m%d%H%M%S%f")[:-3]
        args.task_id = formatted_time

    if(args.language == None):
        args.language = languages[0]

    transcribe = AudioTranscribe()

     # --- 音频智能切片 ---
    total_duration = transcribe.get_audio_duration(args.input)
    if total_duration == 0:
       print("无法处理时长为 0 的音频")
       return

    #manager = ModelManager(MODELS_DIR)
    #if(manager.check_model_exists(f"nvidia/{model_list[args.language]}") == False):
    #    manager.download_model(f"nvidia/{model_list[args.language]}")

    args.model = args.model or model_list[args.language]

    # --- 加载模型 ---
    print(f'加载模型：{args.model}')
    import nemo.collections.asr as nemo_asr # type: ignore
    asr_model = nemo_asr.models.ASRModel.from_pretrained(model_name=f"nvidia/{args.model}")

    # --- 使用 Silero VAD 进行智能切片 ---
    print(f"文件总时长: {total_duration:.2f}s")
    chunk_paths = transcribe.create_intelligent_chunks_silero(
        args.input, 
        total_duration, 
        args.task_id, 
        args.temp_dir or TEMPORARY_DIR, 
        args.chunk_minutes, 
        args.min_silence_duration)
    print(f"[{args.task_id}] 共创建 {len(chunk_paths)} 个切片")
    
    # --- 循环转录并合并结果 ---
    all_segments = transcribe.transcribe_chunk_paths(asr_model, chunk_paths, args.task_id)

    # --- 转录结果转 SRT ---
    print(f"[{args.task_id}] 所有切片转录完成，正在合并结果。")
    srt_result = transcribe.segments_to_srt(all_segments)

    # 打开文件，使用'w'模式（如果文件不存在则创建，存在则覆盖）
    try:
        with open(args.output, 'w', encoding='utf-8') as file:
            file.write(srt_result)
            print(f"[{args.task_id}] srt 写入完成")
    except Exception as e:
        print(f"[{args.task_id}] 写入文件时出错: {e}")
        # 备选方案：清理特殊字符后重试
        import re
        srt_cleaned = re.sub(r'[^\x00-\x7F\u4e00-\u9fff]', '', srt_result)  # 移除非ASCII和中文字符外的字符
        with open(args.output, 'w', encoding='utf-8') as file:
            file.write(srt_cleaned)
            print(f"[{args.task_id}] 已清理特殊字符并写入完成")

    # --- 清理临时文件 ---
    transcribe.clean_temp_chunks(args.task_id)

if __name__ == '__main__':
    transcribe_audio()